import pymysql


class StockDBOperator:
    def __init__(self, host='localhost', port=3306, user='your_username', password='your_password', database='stock_db'):
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database
        self.connection = None
        self.cursor = None
        
    def set_table_name(self,table_name):
        self.table_name = table_name
        

    def connect(self):
        try:
            self.connection = pymysql.connect(
                host=self.host,
                port=self.port,
                user=self.user,
                password=self.password,
                database=self.database,
                charset='utf8mb4',
                cursorclass=pymysql.cursors.DictCursor
            )
            self.cursor = self.connection.cursor()
        except pymysql.Error as e:
            print(f"连接数据库时出错: {e}")

    def disconnect(self):
        if self.cursor:
            self.cursor.close()
        if self.connection:
            self.connection.close()
            
    def create_tables(self,sql):
        try:
            self.connect()                        
            self.cursor.execute(sql)
            self.connection.commit()
            print("创建表成功")
        except pymysql.Error as e:
            self.connection.rollback()
        finally:
            self.disconnect()
            
    def inser_df_data(self,name,code,stock_df):
        ret = True
        try:
            self.connect()
            with self.connection.cursor() as cursor:
                for index, row in stock_df.iterrows():
                    # print(row)
                    # 插入数据的 SQL 语句
                    insert_sql = f"""
                    INSERT INTO {self.table_name} (日期, 股票代码, 名称, 开盘, 收盘, 最高, 最低, 成交量, 成交额, 振幅, 涨跌幅, 涨跌额, 换手率, 量比)
                    VALUES ('{row['日期']}', '{code}', '{name}', {row['开盘']}, {row['收盘']}, {row['最高']}, {row['最低']}, {row['成交额(亿)']}, {row['成交量(万手)']}, {row['振幅']}, {row['涨跌幅']}, {row['涨跌额']}, {row['换手率']}, {row['量比']});
                    """
                    # 执行插入语句
                    # print(insert_sql)
                    cursor.execute(insert_sql)
                    # if index > 1:
                    #     break
            # 提交事务
            self.connection.commit()
            print(f"股票 {name} 的数据插入成功！")
            return ret
        except pymysql.Error as e:
            print(f"插入股票信息时出错: {e}")
            self.connection.rollback()   
            ret = False         
        finally:
            self.disconnect()
            return ret

    def add_stock(self, stock_id, stock_name, listing_date, industry, current_price):
        try:
            self.connect()
            
            sql = "INSERT INTO stocks (stock_id, stock_name, listing_date, industry, current_price) VALUES (%s, %s, %s, %s, %s)"
            self.cursor.execute(sql, (stock_id, stock_name, listing_date, industry, current_price))
            self.connection.commit()
            print("股票信息添加成功")
        except pymysql.Error as e:
            print(f"添加股票信息时出错: {e}")
            self.connection.rollback()
        finally:
            self.disconnect()

    def delete_stock(self, stock_id):
        try:
            self.connect()
            sql = "DELETE FROM stocks WHERE stock_id = %s"
            self.cursor.execute(sql, (stock_id,))
            self.connection.commit()
            print("股票信息删除成功")
        except pymysql.Error as e:
            print(f"删除股票信息时出错: {e}")
            self.connection.rollback()
        finally:
            self.disconnect()

    def update_stock(self, stock_id, stock_name=None, listing_date=None, industry=None, current_price=None):
        try:
            self.connect()
            update_fields = []
            values = []
            if stock_name:
                update_fields.append("stock_name = %s")
                values.append(stock_name)
            if listing_date:
                update_fields.append("listing_date = %s")
                values.append(listing_date)
            if industry:
                update_fields.append("industry = %s")
                values.append(industry)
            if current_price:
                update_fields.append("current_price = %s")
                values.append(current_price)
            values.append(stock_id)

            if update_fields:
                sql = f"UPDATE stocks SET {', '.join(update_fields)} WHERE stock_id = %s"
                self.cursor.execute(sql, tuple(values))
                self.connection.commit()
                print("股票信息更新成功")
            else:
                print("没有需要更新的字段")
        except pymysql.Error as e:
            print(f"更新股票信息时出错: {e}")
            self.connection.rollback()
        finally:
            self.disconnect()

    def get_stock(self, stock_id=None):
        try:
            self.connect()
            if stock_id:
                sql = "SELECT * FROM stocks WHERE stock_id = %s"
                self.cursor.execute(sql, (stock_id,))
            else:
                sql = "SELECT * FROM stocks"
                self.cursor.execute(sql)
            results = self.cursor.fetchall()
            return results
        except pymysql.Error as e:
            print(f"查询股票信息时出错: {e}")
        finally:
            self.disconnect()

    